grammar Strings;

@header {

import java.util.Collections;
import java.util.Set;
import java.util.HashSet;
import java.util.Map;
import java.util.HashMap;
import java.util.List;
import java.util.ArrayList;
import java.util.Iterator;

}

@members {
    private String className = null;

    private Map<String, Integer> strVars = new HashMap<String, Integer>();
    private Map<String, Integer> intVars = new HashMap<String, Integer>();

    private int localVariables = 1;
    private int currentStack = 1;
    private int maxStack = 5;

    private int bufferedReaderIdentifier = -1;

    private void recordPush() {
        ++currentStack;
        maxStack = Math.max(maxStack, currentStack);
    }

    private void recordPop() {
        --currentStack;
    }

    public void setClassName(String name) {
        this.className = name;
    }

    public static class AppendableIterable implements Iterable<String> {
        private final AppendableIterable p1;
        private final AppendableIterable p2;
        private final String my;

        public AppendableIterable() {
            this.p1 = null;
            this.p2 = null;
            this.my = null;
        }

        public AppendableIterable(AppendableIterable previous, String my) {
            this.p1 = previous;
            this.p2 = null;
            this.my = my;
        }

        public AppendableIterable(AppendableIterable p1, AppendableIterable p2) {
            this.p1 = p1;
            this.p2 = p2;
            this.my = null;
        }

        public int size() {
            return (my == null ? 0 : 1) + (p1 == null ? 0 : p1.size()) + (p2 == null ? 0 : p2.size());
        }

        public AppendableIterable append(String s) {
            return new AppendableIterable(this, s);
        }

        public AppendableIterable append(AppendableIterable s) {
            return new AppendableIterable(this, s);
        }

        protected void fillList(List<String> s) {
            if (p1 != null) {
                p1.fillList(s);
            }
            if (p2 != null) {
                p2.fillList(s);
            }
            if (my != null) {
                s.add(my);
            }
        }

        public Iterator<String> iterator() {
            List<String> list = new ArrayList<String>(size());
            fillList(list);
            return list.iterator();
        }
    }
    private static final AppendableIterable Z = new AppendableIterable();
}

program returns [AppendableIterable code]
    : {
          $code = Z;
          bufferedReaderIdentifier = ++localVariables;
      }
      (statement { $code = $code.append($statement.code); })*
      {
          if (currentStack != 1) {
              throw new AssertionError("Something wrong with the stack: " + currentStack);
          }
          AppendableIterable pre = Z.append(".class " + className);
          pre = pre.append(".super java/lang/Object");
          pre = pre.append(".method private <init>()V");
          pre = pre.append("aload_0");
          pre = pre.append("invokenonvirtual java/lang/Object/<init>()V");
          pre = pre.append("return");
          pre = pre.append(".end method");
          pre = pre.append(".method static public main([Ljava/lang/String;)V");
          pre = pre.append(".limit stack " + maxStack);
          pre = pre.append(".limit locals " + (localVariables + 1));
          pre = pre.append("new java/io/BufferedReader");
          pre = pre.append("dup");
          pre = pre.append("new java/io/InputStreamReader");
          pre = pre.append("dup");
          pre = pre.append("getstatic java/lang/System/in Ljava/io/InputStream;");
          pre = pre.append("invokenonvirtual java/io/InputStreamReader/<init>(Ljava/io/InputStream;)V");
          pre = pre.append("invokenonvirtual java/io/BufferedReader/<init>(Ljava/io/Reader;)V");
          pre = pre.append("astore " + bufferedReaderIdentifier);

          $code = pre.append($code);
          $code = $code.append("aload " + bufferedReaderIdentifier);
          $code = $code.append("invokevirtual java/io/Reader/close()V");
          $code = $code.append("return");
          $code = $code.append(".end method");
      }
    ;

statement returns [AppendableIterable code]
    : writeOperation      { $code = $writeOperation.code; }
    | variableDeclaration { $code = $variableDeclaration.code; }
    | variableAssignment  { $code = $variableAssignment.code; }
    ;

/***
 * writeOperation : K_WRITE '(' stringExpression ')'
 *                | K_WRITELINE '(' stringExpression ')'
 *                ;
 */
writeOperation returns [AppendableIterable code]
    : K_WRITE '(' stringExpression ')' ';'
      {
          $code = $stringExpression.code;
          $code = $code.append("getstatic java/lang/System/out Ljava/io/PrintStream;");
          recordPush();
          $code = $code.append("swap");
          $code = $code.append("invokevirtual java/io/PrintStream/print(Ljava/lang/String;)V");
          recordPop();
          recordPop();
      }
    | K_WRITELINE '(' stringExpression ')' ';'
      {
          $code = $stringExpression.code;
          $code = $code.append("getstatic java/lang/System/out Ljava/io/PrintStream;");
          recordPush();
          $code = $code.append("swap");
          $code = $code.append("invokevirtual java/io/PrintStream/println(Ljava/lang/String;)V");
          recordPop();
          recordPop();
      }
    ;

/***
 * variableDeclaration : K_VAR ID '=' (stringExpression | intExpression) ';'
 *                     ;
 */
variableDeclaration returns [AppendableIterable code]
    : K_VAR ID '=' (
      (stringExpression)=> se=stringExpression
        {
            strVars.put($ID.text, ++localVariables);
            $code = ($se.code).append("astore " + localVariables);
            recordPop();
        }
      | intExpression
        {
            intVars.put($ID.text, ++localVariables);
            $code = ($intExpression.code).append("istore " + localVariables);
            recordPop();
        }
      ) ';'
    ;

/***
 * variableAssignment : {strVars.containsKey(input.LT(1).getText())}?=> ID '=' stringExpression ';'
 *                    | {intVars.containsKey(input.LT(1).getText())}?=> ID '=' intExpression ';'
 *                    ;
 */
variableAssignment returns [AppendableIterable code]
    : {strVars.containsKey(input.LT(1).getText())}?=> ID '=' stringExpression ';'
      {
          $code = $stringExpression.code.append("astore " + strVars.get($ID.text));
          recordPop();
      }
    | {intVars.containsKey(input.LT(1).getText())}?=> ID '=' intExpression ';'
      {
          $code = $intExpression.code.append("istore " + intVars.get($ID.text));
          recordPop();
      }
    ;

/***
 * stringExpression : stringValue ('+' stringValue)*
 *                  ;
 */
stringExpression returns [AppendableIterable code]
@init {
    boolean hasConcat = false;
}
    : s1=stringValue
      {
          $code = $s1.code;
      }
      ('+' s2=stringValue
          {
              if (!hasConcat) {
                  $code = $code.append("new java/lang/StringBuilder");
                  recordPush();
                  $code = $code.append("dup");
                  recordPush();
                  $code = $code.append("invokenonvirtual java/lang/StringBuilder/<init>()V");
                  recordPop();
                  $code = $code.append("swap");
                  $code = $code.append("invokevirtual java/lang/StringBuilder/append(Ljava/lang/String;)Ljava/lang/StringBuilder;");
                  recordPop();
                  hasConcat = true;
              }
              $code = $code.append($s2.code);
              $code = $code.append("invokevirtual java/lang/StringBuilder/append(Ljava/lang/String;)Ljava/lang/StringBuilder;");
              recordPop();
          }
      )*
      {
          if (hasConcat) {
              $code = $code.append("invokevirtual java/lang/StringBuilder/toString()Ljava/lang/String;");
          }
      }
    ;

/***
 * stringValue : STRING
 *             | {strVars.containsKey(input.LT(1).getText())}?=> ID
 *             | '(' stringExpression ')'
 *             | K_TOSTR '(' intExpression ')'
 *             | K_NEXTLINE
 *             ;
 */
stringValue returns [AppendableIterable code]
    : STRING
      {
          $code = Z.append("ldc " + $STRING.text);
          recordPush();
      }
    | {strVars.containsKey(input.LT(1).getText())}?=> ID
      {
          $code = Z.append("aload " + strVars.get($ID.text));
          recordPush();
      }
    | '(' stringExpression ')'
      {
          $code = $stringExpression.code;
      }
    | K_TOSTR '(' intExpression ')'
      {
          $code = $intExpression.code.append("invokestatic java/lang/String/valueOf(I)Ljava/lang/String;");
      }
    | K_READLINE
      {
          $code = Z.append("aload " + bufferedReaderIdentifier);
          recordPush();
          $code = $code.append("invokevirtual java/io/BufferedReader/readLine()Ljava/lang/String;");
      }
    ;

/***
 * intExpression : intMultExpression (('+' intMultExpression) | ('-' intMultExpression))*
 *               ;
 */
intExpression returns [AppendableIterable code]
    : i1=intMultExpression
      {
          $code = $i1.code;
      }
      (('+' i2=intMultExpression
          {
              $code = $code.append($i2.code).append("iadd");
              recordPop();
          }
      ) | ('-' i3=intMultExpression
          {
              $code = $code.append($i3.code).append("isub");
              recordPop();
          }
      ))*
    ;

/***
 * intMultExpression : intValue (('*' intValue) | ('/' intValue))*
 *                   ;
 */
intMultExpression returns [AppendableIterable code]
    : i1=intValue
      {
          $code = $i1.code;
      }
      (('*' i2=intValue
          {
              $code = $code.append($i2.code).append("imul");
              recordPop();
          }
      ) | ('/' i3=intValue
          {
              $code = $code.append($i3.code).append("idiv");
              recordPop();
          }
      ))*
    ;

/***
 * intValue : INT
 *          | {intVars.containsKey(input.LT(1).getText())}?=> ID
 *          | '(' intExpression ')'
 *          | K_LEN '(' stringExpression ')'
 *          | K_TOINT '(' stringExpression ')'
 *          ;
 */
intValue returns [AppendableIterable code]
    : INT
      {
          try {
              $code = Z.append("ldc " + Integer.parseInt($INT.text));
              recordPush();
          } catch (NumberFormatException ex) {
              throw ex;
          }
      }
    | {intVars.containsKey(input.LT(1).getText())}?=> ID
      {
          $code = Z.append("iload " + intVars.get($ID.text));
          recordPush();
      }
    | '(' intExpression ')'
      {
          $code = $intExpression.code;
      }
    | K_LEN '(' stringExpression ')'
      {
          $code = $stringExpression.code.append("invokevirtual java/lang/String/length()I");
      }
    | K_TOINT '(' stringExpression ')'
      {
          $code = $stringExpression.code.append("invokestatic java/lang/Integer/parseInt(Ljava/lang/String;)I");
      }
    ;

K_VAR       : 'var';
K_IF        : 'if';
K_ELSE      : 'else';
K_FOR       : 'for';
K_LEN       : 'length';
K_TOINT     : 'toInt';
K_TOSTR     : 'toString';
K_READLINE  : 'readLine';
K_WRITE     : 'write';
K_WRITELINE : 'writeLine';

ID      : ('a'..'z' | 'A'..'Z')+;
INT     : '-'? ('0'..'9')+;
STRING  : '"' (~('"' | '\n'))* '"';
WS      : (' ' | '\t' | '\n' | '\r')+ { skip(); };
